#ifndef MACDONALDFOUR_HPP
#define MACDONALDFOUR_HPP

#include "edgeTraversal.hpp"
#include "sweSolver.hpp"
#include "solverBase.hpp"
#include "meshTypes.hpp"
#include "meshTypeFunctions.hpp"
#include "magic_enum.hpp"
#include "unifiedMathFunctions.hpp"
#include <functional>
#include <vector>
#include <cmath>

using namespace hms;

class LegendrePolynomial {
public:
    LegendrePolynomial(double lowerBound, double upperBound, size_t numberOfIterations)
        : mLowerBound(lowerBound), mUpperBound(upperBound), mNumberOfIterations(numberOfIterations), mWeight(numberOfIterations+1), mRoot(numberOfIterations+1) {
        calculateWeightAndRoot();
    }

    const std::vector<double> & getWeight() const {
        return mWeight;
    }

    const std::vector<double> & getRoot() const {
        return mRoot;
    }

private:
    const static double EPSILON;

    struct Result {
        double value;
        double derivative;

        Result() : value(0), derivative(0) {}
        Result(double val, double deriv) : value(val), derivative(deriv) {}
    };

    void calculateWeightAndRoot() {
        for(int step = 0; step <= mNumberOfIterations; step++) {
            double root = cos(M_PI * (step-0.25)/(mNumberOfIterations+0.5));
            Result result = calculatePolynomialValueAndDerivative(root);

            double newtonRaphsonRatio;
            do {
                newtonRaphsonRatio = result.value/result.derivative;
                root -= newtonRaphsonRatio;
                result = calculatePolynomialValueAndDerivative(root);
            } while (fabs(newtonRaphsonRatio) > EPSILON);

            mRoot[step] = root;
            mWeight[step] = 2.0/((1-root*root)*result.derivative*result.derivative);
        }
    }

    Result calculatePolynomialValueAndDerivative(double x) {
        Result result(x, 0);

        double value_minus_1 = 1;
        const double f = 1/(x*x-1);
        for(int step = 2; step <= mNumberOfIterations; step++) {
            const double value = ((2*step-1)*x*result.value-(step-1)*value_minus_1)/step;
            result.derivative = step*f*(x*value - result.value);

            value_minus_1 = result.value;
            result.value = value;
        }

        return result;
    }

    const double mLowerBound;
    const double mUpperBound;
    const int mNumberOfIterations;
    std::vector<double> mWeight;
    std::vector<double> mRoot;
};

const double LegendrePolynomial::EPSILON = 1e-15;

double gaussLegendreIntegral(double a, double b, int n, const std::function<double (double)> &f) {
    const LegendrePolynomial legendrePolynomial(a, b, n);
    const std::vector<double> & weight = legendrePolynomial.getWeight();
    const std::vector<double> & root = legendrePolynomial.getRoot();

    const double width = 0.5*(b-a);
    const double mean = 0.5*(a+b);

    double gaussLegendre = 0;
    for(int step = 1; step <= n; step++) {
        gaussLegendre += weight[step]*f(width * root[step] + mean);
    }

    return gaussLegendre * width;
}

struct macDonaldFour{
    

	hms::EdgeTraversal traversal;
	hms::LimiterFunction limFunc;
	bool useMuscl, useFriction;

	

	template<typename MeshType>
	void runCase(Index nx, Index ny){
		static_assert( std::is_same_v<MeshType, UniMesh> || std:: is_same_v<MeshType, RectMesh>, 
		"1D Dam Break Verification case can only run on MeshTypes: RectMesh, UniMesh. The entered type does not match. Please choose on of the two options.");

		using namespace hms;

		std::cout << "Verification test case: MacDonald's test problem #4 (1D Hydraulic Jump)" << std::endl;

		/* Domain: 1000 m x 10 m*/
		scalar l { 1000.0 }; // channel length
		scalar w { 10.0 };   // channel width
		scalar discharge { 20.0 };
		scalar n {0.02};     // Manning's n
		Index gaussIterations {1000}; // for gauss legendre integration

		/* mesh */
		scalar dx {l / nx};
		scalar dy {w / ny};
		Index  ng {2};

		auto mesh = uniMeshAs<MeshType>(nx,ny,dx,dy,ng);

		mesh.printInfo();

		

		SWESolver<MeshType> solver;
		solver.edgeTraversal(traversal);

		std::cout << "Edge traversal: "
			<< magic_enum::enum_name( solver.edgeTraversal() )
			<< "\nUsing second order reconstruction: " << useMuscl
			<< "\n";

		solver.useFriction(useFriction);
		solver.useSourceFields(false);
		solver.useMuscl(useMuscl);
		solver.limiterFunction(limFunc);

		solver.v() = Field{mesh, 2};
		solver.d() = Field{mesh, 1};
		solver.z() = Field{mesh, 1};
		solver.v().values() = 0;
		solver.z().values() = 1.6;

		solver.n() = ConstantField<Uniform, MeshType>{ mesh, ArrayNs<1>{0.020} };
		solver.r() = SourceField<MeshType,   Uniform>{ mesh, ArrayNs<1>{0    } };
		solver.i() = SourceField<MeshType,   Uniform>{ mesh, ArrayNs<1>{0    } };

		/* water depth derivative */
		auto depthDerivative = [&](scalar x){
		if( x <= 300.) 
			return -0.723499e-3 * square(1.0 / cosh( x/l - 3./10. ));
		
		else if(300. < x && x <= 600.) 
			return -0.723449e-3 * square(1.0 / cosh( 6. * ( x/l - 3./10. ) ) );
		
		else 
			return -1./50. * (1. * -0.111051 * exp(-20. * 1. * ( x/l - 3./5.)) +
		        2. *  0.026876 * exp(-20. * 2. * ( x/l - 3./5.)) +
				3. * -0.215767 * exp(-20. * 3. * ( x/l - 3./5.))) 
				+ 3./5000. * exp( x/l -1.);
	};

		/* water depth */
		auto depth = [&](scalar x){
			if( x <= 300.0 ) 
				return 0.723449 * (1. - 1./6. * tanh(6. * ( x/l - 3./10. ) ) );
			
			else if(300. < x && x <= 600.) 
				return 0.723449 * (1. - 1./6. * tanh(6. * ( x/l - 3./10. ) ) );
			
			else
				return 3. / 4. + -0.111051 * exp(-20. * 1. * ( x/l - 3./5.)) + 
					0.026876 * exp(-20. * 2. * ( x/l - 3./5.)) + 
					-0.217567 * exp(-20. * 3. * ( x/l - 3./5.)) + 
					3./5. * exp( x/l - 1.);
		};
				
		/* required functions for bottom slope */
		auto f1 = [&](scalar d){
			double result {1. - square(discharge) * w / pow((phys::g * w * d), 3)};		
			return result;
		};

		auto f2 = [&](scalar d){
			double result {square(discharge) * square(n) * pow(w, 4./3.) 
				/ pow(w*d , 10./3.)};
			return result;
		};

		/* bottom slope */
		auto bottomSlope = [&](scalar x){
			scalar d { depth(x) };

			return f1(d) * depthDerivative(x) + f2(d);
		};
		
		/* reshaped functions */
		auto getReshaped = [&](){
			return std::pair{
				solver.z().domainValues().reshaped(ny,nx),
				solver.d().domainValues().reshaped(ny,nx)
			};
		};
		std::cout << nx << "   " << ny << std::endl;

		

		auto [dReshaped, zReshaped] = getReshaped();

		OMP_PRAGMA( parallel for )
		for(Index i = 0; i< nx; i++) {
			std::cout << i << std::endl;

			zReshaped.col(i) = gaussLegendreIntegral(mesh.cell(i,0).centroid()[0], l, gaussIterations, bottomSlope);
			
			
			dReshaped.col(i) = 2.0;
			
		}
		std::cout << "TILLHERE" << std::endl;
		

		/* calculate bottom elevation and water depth for each cell */
		// mesh.forEachCell( [&](const auto& cell){
		// 	solver.z().values()(0, cell.index()) = gaussLegendreIntegral(cell.centroid()[0], 1000, 1000, bottomSlope);
		// 	solver.d().values()(0, cell.index()) = 2.0;
		// });

		solver.initialise();

		Time time {10};

		Timer t;
		std::cout << std::setprecision(5) << std::fixed;
		std::cout << "Time: " << time.current() << "s";
		while ( !time.atEnd() ){
			solver.advance( time );
			std::cout << "\b\b\b\b\b\b\b\b";
			std::cout << time.current() << "s";
		}
		double elapsed { t.elapsed() };
		double execTimeSum {0};
		for ( auto execTime : solver.execTimings() )
			execTimeSum += execTime;

		std::cout
			<< "\nComputation time: " << elapsed << "s"
			<< "\nSum of execution times logged by solver: " << execTimeSum << "s"
			<< "\n\t" << 100 * solver.execTimings()[0] / execTimeSum
				<< "% for time step calculation"
				<< " (" << solver.execTimings()[0] << "s)"
			<< "\n\t" << 100 * solver.execTimings()[1] / execTimeSum
				<< "% for fluxes (including slope)"
				<< " (" << solver.execTimings()[1] << "s)"
			<< "\n\t" << 100 * solver.execTimings()[2] / execTimeSum
				<< "% for sources " << ( useFriction ?
					std::string{ "(including friction)" } :
					std::string{ "(friction disabled)" } )
				<< " (" << solver.execTimings()[2] << "s)"
			<< "\n\t" << 100 * solver.execTimings()[3] / execTimeSum
				<< "% for new state"
				<< " (" << solver.execTimings()[3] << "s)"
			<< "\n\t" << 100 * solver.execTimings()[4] / execTimeSum
				<< "% for updating boundaries"
				<< " (" << solver.execTimings()[4] << "s)"
			<< "\n";
	}
};


#endif